#!/usr/bin/env python3
"""
Build LWG training dataset from exported per-sequence CSVs and SMD_DanceTrack GT.

Labeling rule (MVP):
 - Match current sample to a GT ID by IoU > iou_thresh at time t.
 - If matched ID exists in GT for >= M frames within [t+1, t+K], label 1 (should write), else 0.

Output: a single CSV with selected features and label y.
Features used (must match runtime gating): [s_det, sim1, margin, area_ratio, aspect_ratio]
"""
import sys
import os
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../../")))
import os
import re
import csv
import math
import argparse
from collections import defaultdict

from data.joint_dataset import dataset_classes


def iou_xywh(b1, b2):
    x1, y1, w1, h1 = b1
    x2, y2, w2, h2 = b2
    ax1, ay1, ax2, ay2 = x1, y1, x1 + w1, y1 + h1
    bx1, by1, bx2, by2 = x2, y2, x2 + w2, y2 + h2
    ix1, iy1 = max(ax1, bx1), max(ay1, by1)
    ix2, iy2 = min(ax2, bx2), min(ay2, by2)
    iw, ih = max(0.0, ix2 - ix1), max(0.0, iy2 - iy1)
    inter = iw * ih
    if inter <= 0:
        return 0.0
    area_a = w1 * h1
    area_b = w2 * h2
    union = area_a + area_b - inter + 1e-6
    return inter / union


def load_gt(seq_dir):
    gt_path = os.path.join(seq_dir, "gt", "gt.txt")
    frame_to_objs = defaultdict(list)
    if not os.path.exists(gt_path):
        return frame_to_objs
    with open(gt_path, "r") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            fr, tid, x, y, w, h, conf, _, _ = line.split(",")
            fr = int(float(fr))
            tid = int(float(tid))
            x = float(x)
            y = float(y)
            w = float(w)
            h = float(h)
            frame_to_objs[fr].append((tid, (x, y, w, h)))
    return frame_to_objs


def parse_row(row):
    def fnum(v, default=0.0):
        try:
            if v == "" or v is None:
                return default
            return float(v)
        except Exception:
            return default

    def fint(v, default=0):
        try:
            if v == "" or v is None:
                return default
            return int(float(v))
        except Exception:
            return default

    return {
        "seq": row[0],
        "t": fint(row[1]),
        "idx": fint(row[2]),
        "id_label": fint(row[3], -1),
        "track_id": fint(row[4], -1),
        "s_det": fnum(row[5]),
        "sim1": fnum(row[6]),
        "sim2": fnum(row[7]),
        "margin": fnum(row[8]),
        "track_age": fint(row[9]),
        "delta_t": fint(row[10]),
        "area_ratio": fnum(row[11]),
        "aspect_ratio": fnum(row[12]),
        "x": fnum(row[13]),
        "y": fnum(row[14]),
        "w": fnum(row[15]),
        "h": fnum(row[16]),
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dump-dir", type=str, required=True, help="Directory of per-sequence CSV dumps")
    ap.add_argument("--data-root", type=str, required=True, help="Dataset root (contains SMD_DanceTrack)")
    ap.add_argument("--split", type=str, default="train")
    ap.add_argument("--out-path", type=str, required=True)
    ap.add_argument("--K", type=int, default=20)
    ap.add_argument("--M", type=int, default=8)
    ap.add_argument("--iou-thresh", type=float, default=0.5)
    ap.add_argument(
        "--dataset",
        type=str,
        default="SMD_DanceTrack_Spilt",
        help="Dataset key registered in dataset_classes (e.g., DanceTrack, MaritimeTrack_Full_20250413)."
    )
    args = ap.parse_args()

    dataset_name = args.dataset
    if dataset_name not in dataset_classes:
        raise ValueError(f"Unknown dataset '{dataset_name}'. Available: {list(dataset_classes.keys())}")

    # 加载真实数据集序列信息
    dataset = dataset_classes[dataset_name](
        data_root=args.data_root,
        split=args.split,
        load_annotation=True
    )
    seq_to_gt = {}
    dataset_root_dir = getattr(dataset, "data_dir", os.path.join(args.data_root, dataset_name))
    for seq_name, seq_info in dataset.sequence_infos.items():
        # 如果没有 seq_dir，则自动生成
        seq_dir = seq_info.get(
            "seq_dir",
            os.path.join(dataset_root_dir, args.split, seq_name)
        )
        seq_to_gt[seq_name] = load_gt(seq_dir)

    os.makedirs(os.path.dirname(args.out_path), exist_ok=True)
    with open(args.out_path, "w", newline="") as fout:
        writer = csv.writer(fout)
        header = [
            "seq", "t", "idx", "id_label", "track_id",
            "s_det", "sim1", "margin", "area_ratio", "aspect_ratio",
            "best_iou", "delta_sim1", "y"
        ]
        writer.writerow(header)

        dump_files = [f for f in os.listdir(args.dump_dir) if f.endswith(".txt") or f.endswith(".csv")]

        rank_suffix_pattern = re.compile(r"_rank\d+$")
        for csv_name in dump_files:
            seq_raw = os.path.splitext(csv_name)[0]
            seq_lookup = rank_suffix_pattern.sub("", seq_raw)
            gt = seq_to_gt.get(seq_lookup, None)
            if gt is None:
                print(f"[WARN] GT not found for sequence {seq_raw}")
                continue
            dump_path = os.path.join(args.dump_dir, csv_name)
            samples = []
            with open(dump_path, "r") as fin:
                reader = csv.reader(fin)
                for row in reader:
                    if not row:
                        continue
                    r = parse_row(row)
                    fr = r["t"] + 1
                    cur_box = (r["x"], r["y"], r["w"], r["h"])
                    cur_tid = -1
                    best_iou = 0.0
                    for tid, gtb in gt.get(fr, []):
                        iou = iou_xywh(cur_box, gtb)
                        if iou > best_iou:
                            best_iou = iou
                            cur_tid = tid
                    sample = {
                        "seq": seq_lookup,
                        "t": r["t"],
                        "idx": r["idx"],
                        "id_label": r["id_label"],
                        "track_id": r["track_id"],
                        "s_det": r["s_det"],
                        "sim1": r["sim1"],
                        "margin": r["margin"],
                        "area_ratio": r["area_ratio"],
                        "aspect_ratio": r["aspect_ratio"],
                        "box": cur_box,
                        "best_iou": 0.0,
                        "delta_sim1": 0.0,
                        "gt_tid": cur_tid if best_iou >= args.iou_thresh else -1,
                        "label": 0,
                    }
                    samples.append(sample)

            # compute labels
            tid_to_indices = defaultdict(list)
            for idx, sample in enumerate(samples):
                cur_tid = sample["gt_tid"]
                if cur_tid >= 0:
                    cnt = 0
                    fr = sample["t"] + 1
                    for dt in range(1, args.K + 1):
                        fr2 = fr + dt
                        if any(tid == cur_tid for tid, _ in gt.get(fr2, [])):
                            cnt += 1
                    sample["label"] = 1 if cnt >= args.M else 0
                    tid_to_indices[cur_tid].append(idx)
                else:
                    sample["label"] = 0

            # compute temporal consistency features (best IoU with previous box, similarity delta)
            samples.sort(key=lambda s: (s["t"], s["idx"]))
            prev_by_track = {}
            for sample in samples:
                tid = sample["track_id"]
                box = sample.get("box", None)
                sim1_val = sample.get("sim1", 0.0)
                best_iou_track = 0.0
                delta_sim = 0.0
                if (
                    box is not None
                    and tid is not None
                    and tid >= 0
                    and tid in prev_by_track
                ):
                    prev = prev_by_track[tid]
                    best_iou_track = iou_xywh(box, prev["box"])
                    delta_sim = sim1_val - prev["sim1"]
                if not math.isfinite(best_iou_track):
                    best_iou_track = 0.0
                best_iou_track = max(0.0, min(1.0, best_iou_track))
                if not math.isfinite(delta_sim):
                    delta_sim = 0.0
                sample["best_iou"] = best_iou_track
                sample["delta_sim1"] = delta_sim
                if box is not None and tid is not None and tid >= 0:
                    prev_by_track[tid] = {"box": box, "sim1": sim1_val}

            for sample in samples:
                writer.writerow([
                    sample["seq"], sample["t"], sample["idx"], sample["id_label"], sample["track_id"],
                    f"{sample['s_det']:.6f}",
                    f"{sample['sim1']:.6f}",
                    f"{sample['margin']:.6f}",
                    f"{sample['area_ratio']:.6f}",
                    f"{sample['aspect_ratio']:.6f}",
                    f"{sample['best_iou']:.6f}",
                    f"{sample['delta_sim1']:.6f}",
                    sample["label"],
                ])

    print(f"Dataset written to {args.out_path}")


if __name__ == "__main__":
    main()
